/*
 * File:    SecureChannel.java
 * Created: 17-Oct-2005
 * Version: $Id: //depot/bitgate/bsse/src/bitgate/util/socket/SecureChannel.java#1 $
 *
 * COPYRIGHT (C) 2005, Bitgate Software, LLC.  All Rights Reserved.
 *
 * software@bitgatesoftware.com
 */

package com.bitgate.util.socket;

import static com.bitgate.util.debug.Debug.*;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.ClosedChannelException;
import java.nio.channels.SocketChannel;

import javax.net.ssl.SSLEngine;
import javax.net.ssl.SSLEngineResult;
import static javax.net.ssl.SSLEngineResult.Status.*;
import static javax.net.ssl.SSLEngineResult.HandshakeStatus.*;
import javax.net.ssl.SSLException;
import javax.net.ssl.SSLSession;

/**
 * This is a secure non-blocking socket class using an <code>SSLEngine</code> implementation.  It uses
 * the standard TLS encryption unless otherwise specified in the configuration file.
 *
 * @author Kenji Hollis &lt;kenji@bitgatesoftware.com&gt;
 * @version $Id: //depot/bitgate/bsse/src/bitgate/util/socket/SecureChannel.java#1 $
 */
public class SecureChannel extends AbstractClientChannel
{
    private final SSLEngine engine;
    private final ByteBuffer inAppData, outAppData;
    private final ByteBuffer inNetData, outNetData;
    private WriteFuture writeFuture;
    private byte[] secureWriteData;
    
    private boolean closed = false;

    public SecureChannel(SocketChannel wrappedChannel, SSLEngine engine)
    {
    	super(wrappedChannel);
        this.engine = engine;
        this.writeData = null;

        SSLSession session = engine.getSession();

        this.inAppData = ByteBuffer.allocate(session.getApplicationBufferSize());
        this.outAppData = ByteBuffer.allocate(session.getApplicationBufferSize());
        this.inNetData = ByteBuffer.allocate(session.getPacketBufferSize());
        this.outNetData = ByteBuffer.allocate(session.getPacketBufferSize());
        
        this.writeFuture = new AbstractWriteFuture();
    }

    @Override
    public void close() throws IOException
    {
        if (!this.closed) {
            try {
                this.closed = true;
                this.engine.closeOutbound();
                sslLoop(wrap());
                this.wrappedChannel.close();
            } finally {
                this.closed = true;
            }
        }
        
        secureWriteData = null;
    }

    @Override
    public int read(ByteBuffer byteBuffer) throws IOException, ClosedChannelException
    {
    	if (!isOpen()) {
    		throw new ClosedChannelException();
    	}
    	
        if (isOpen()) {
            try {
                sslLoop(unwrap());
            } catch(SSLException e) {
                throw new IOException("SSLException: " + e.getMessage());
            } catch(ClosedChannelException e) {
                close();
            }
        }

        this.inAppData.flip();
        int posBefore = this.inAppData.position();
        byteBuffer.put(this.inAppData);
        int posAfter = this.inAppData.position();
        this.inAppData.compact();

        SocketStatistics.getDefault().increaseReads((posAfter - posBefore));
        
        if (posAfter - posBefore > 0) {
            return (posAfter - posBefore);
        }

        if (isOpen()) {
            return 0;
        }
        
        return -1;
    }

    @Override
    public ClientChannel write(byte[] data) throws ClosedChannelException
    {
    	if (!isOpen()) {
    		throw new ClosedChannelException();
    	}
    	
        if (!isOpen()) {
            return this;
        }

    	if (secureWriteData != null) {
        	byte[] outData = new byte[secureWriteData.length + data.length];
        	
        	System.arraycopy(secureWriteData, 0, outData, 0, secureWriteData.length);
        	System.arraycopy(data, 0, outData, secureWriteData.length, data.length);
    		
        	secureWriteData = outData;
        	outData = null;
    	} else {
    		secureWriteData = data;
    	}

        return this; 
    }
    
    // Needs to be fixed!!
    @Override
    public boolean signalWrite() throws IOException
    {
    	if (secureWriteData == null) {
    		return false;
    	}
    	
    	int bytesLeft = secureWriteData.length;
    	ByteBuffer byteBuffer = ByteBuffer.wrap(secureWriteData);
    	int posBefore, posAfter;

    	posBefore = byteBuffer.position();
    	
        if (byteBuffer.remaining() < this.outAppData.remaining()) {
            this.outAppData.put(byteBuffer);
        } else {
            while(byteBuffer.hasRemaining() && this.outAppData.hasRemaining()) {
                this.outAppData.put(byteBuffer.get());
            }
        }

        posAfter = byteBuffer.position();
        
        if (isOpen()) {
            try {
                while(true) {
                    SSLEngineResult r = sslLoop(wrap());

                    if (r.bytesConsumed() == 0 && r.bytesProduced() == 0) {
                        break;
                    }
                }
            } catch(IOException e) {
            	info("Unable to wrap data for sending: " + e.getMessage());
            }
        }

        SocketStatistics.getDefault().increaseWrites(posAfter - posBefore);
        byteBuffer = null;
        
        writeFuture.increaseByteCount(posAfter - posBefore);

        if ((posAfter - posBefore) == bytesLeft) {
    		if (isDebugEnabled()) {
    			debug("Clear data.");
    		}
        	
        	outAppData.clear();
        	outNetData.clear();
        	
        	secureWriteData = null;
    			
			// Calls the "onComplete" function of the write future after writing completes.
			writeFuture.onComplete();
			writeFuture.clearByteCount();
			
			return false;
        } 
        
    	byte[] newData = new byte[(posAfter - posBefore)];
    	
    	System.arraycopy(secureWriteData, (posAfter - posBefore), newData, 0, bytesLeft - (posAfter - posBefore));
    	
    	secureWriteData = newData;
        
        return ((posAfter - posBefore) != bytesLeft);
    }
    
    @Override
    public void setWriteFuture(WriteFuture wFuture)
    {
    	writeFuture = wFuture;
    }
    
    @Override
    public boolean isSecure()
    {
        return true;
    }

    private SSLEngineResult unwrap() throws IOException, SSLException
    {
        int wrapRead = 0;

        while((wrapRead = this.wrappedChannel.read(this.inNetData)) > 0) {
        	if (isDebugEnabled()) {
        		debug("SSLRead: Length read = " + wrapRead);
        	}
        }

        this.inNetData.flip();
        SSLEngineResult ser = this.engine.unwrap(this.inNetData, this.inAppData);
        this.inNetData.compact();

        if (ser.getStatus() == CLOSED || wrapRead == -1) {
            close();
        }

        return ser;
    }

    private SSLEngineResult wrap() throws IOException, SSLException
    {
        SSLEngineResult ser = null;

        this.outAppData.flip();
        ser = this.engine.wrap(this.outAppData, this.outNetData);
        this.outAppData.compact();

        this.outNetData.flip();
        while(this.outNetData.hasRemaining()) {
            this.wrappedChannel.write(this.outNetData);
        }
        this.outNetData.compact();

        return ser;
    }

    private SSLEngineResult sslLoop(SSLEngineResult ser) throws SSLException, IOException
    {
        if (ser == null) {
            return ser;
        }
        
        SSLEngineResult sslEngineResult = ser;

        while(sslEngineResult.getHandshakeStatus() != FINISHED && sslEngineResult.getHandshakeStatus() != NOT_HANDSHAKING &&
        	  !this.closed) {
            switch(sslEngineResult.getHandshakeStatus()) {
                case NEED_TASK:
                	doTasks();
                	sslEngineResult = wrap();
                    break;

                case NEED_WRAP:
                    sslEngineResult = wrap();
                    break;

                case NEED_UNWRAP:
                    sslEngineResult = unwrap();
                    break;
                    
                case NOT_HANDSHAKING:
                	if (isDebugEnabled()) {
                		debug("Not handshaking received, data out of sync, or client did not negotiate data connection.");
                	}
                	break;
                	
                case FINISHED:
                	if (isDebugEnabled()) {
                		debug("Handshake finished.");
                	}
                	break;
                	
                default:
            		if (isDebugEnabled()) {
            			debug("Unknown or unhandled handshake: " + sslEngineResult);
            		}
                	break;
            }
        }

        if (!this.closed) {
            switch(sslEngineResult.getStatus()) {
                case BUFFER_UNDERFLOW:
                case BUFFER_OVERFLOW:
                case CLOSED:
                	if (isDebugEnabled()) {
                		debug("SSLEngine operations done: Closing the socket.");
                	}
                	
                    try {
                        wrappedChannel.close();
                    } finally {
                        closed = true;
                    }
    
                    break;
                
                case OK:
                	if (isDebugEnabled()) {
                		debug("OK status received before closed.");
                	}
                	
                	break;
            }
        }

        return sslEngineResult;
    }
    
    private void doTasks()
    {
        Runnable task;

        while((task = engine.getDelegatedTask()) != null) {
            task.run();
        }
    }
}
